"""

Visualisation and reporting methods


Copyright 2009 Michael Seiler
Rutgers University
miseiler@gmail.com

This file is part of ConsensusCluster.

ConsensusCluster is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

ConsensusCluster is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with ConsensusCluster.  If not, see <http://www.gnu.org/licenses/>.


"""

import logging, sys, os, numpy

from itertools import cycle
from ioutils   import new_filename

class Hmap(object):
    """

    Hmap

        Builds a heatmap of any iterable matrix.  Does nothing if PIL cannot be imported

        Usage:

            Hmap(matrix, bsize = 10)

            matrix      -   Iterable matrix.  Negative values make green squares, positive ones make red ones.
                            Lightness is RELATIVE TO 1.
                            EXPECTS -1 to 1 FRACTIONS IN EACH SQUARE. DOES NOT SCALE ANY LONGER
            bsize       -   The box size for each value in the heatmap
        
        Methods

            save(filename = None)
                Saves the current image to filename in PNG format.  Otherwise filename is 'heatmap.png'. Auto-appends .png to filename.

            show()
                Creates a GTK window displaying the current image
                
                WARNING: Requires the Gtk_UI interface!

        Properties

            im
                The current image

    """

    def __init__(self, matrix, bsize = 10, yellowblue = False):

        self.im = None
        self.bsize = bsize

        self.im = self._create_hmap(matrix, yellowblue)

    def _create_hmap(self, matrix, yellowblue = False):

        from PIL import Image, ImageDraw
       
        size = (len(matrix[0]) * self.bsize, len(matrix) * self.bsize)
    
        im = Image.new('RGBA', size, 'white')
        draw = ImageDraw.Draw(im)
        
        for row in xrange(len(matrix)):
            for col in xrange(len(matrix[row])):
                
                rev   = matrix[row][col] < 0
                value = int(abs(matrix[row][col]) * 255)

                if yellowblue:
                    if rev:
                        colour = (0, 0, value)
                    else:
                        colour = (value, value, 0)

                else:
                    if rev:
                        colour = (0, value, 0)
                    else:
                        colour = (value, 0, 0)

                col_size = col * self.bsize
                row_size = row * self.bsize

                bcol_size = self.bsize + col_size
                brow_size = self.bsize + row_size

                draw.polygon([(col_size, row_size),
                              (bcol_size, row_size),
                              (bcol_size, brow_size),
                              (col_size, brow_size)], outline='black', fill=colour)
            
        return im
    
    def _image_to_pixbuf(self, im):  
        """
        Creates a pixbuf from an Image object.  Probably unnecessary.  In the future I could just invoke a save
        and load then png instead.
        
        """

        import gtk, StringIO
        
        file1 = StringIO.StringIO()  
        im.save(file1, "ppm")  
        contents = file1.getvalue()  
        file1.close()  
        loader = gtk.gdk.PixbufLoader("pnm")  
        loader.write(contents, len(contents))  
        pixbuf = loader.get_pixbuf()  
        loader.close()  
        return pixbuf  
    
    def save(self, filename = None, dpi=None):
        """Save the current image to filename.  Automatically appends a .png extension"""
        
        if dpi is not None:
            print('Setting image dpi to %s' % dpi)
            dpi = (dpi, dpi)
        else:
            #dpi = self.im.info['dpi']
            dpi=None

        self.im.save(new_filename(filename, 'heatmap', '.png'), 'png', dpi=dpi)

    def show(self):
        """Opens a GTK window and puts the heatmap in it.  Intelligent enough to work with the GUI as well."""
        
        window_only = 1 #What needs to be destroyed when the window is destroyed?

        import gtk

        def destroy():

            if window_only:
                window.destroy()
            else:
                gtk.main_quit()

        gtk.gdk.threads_enter()
        window = gtk.Window()
        window.set_title("Showing heatmap...")
        window.set_border_width(10)
        window.set_resizable(False)
        window.connect("delete_event", lambda w, e: destroy())
    
        backbone = gtk.HBox(True)
        image = gtk.Image()
        image.set_from_pixbuf(self._image_to_pixbuf(self.im))
        backbone.pack_start(image)
    
        window.add(backbone)
        window.show_all()
        gtk.gdk.threads_leave()

        if gtk.main_level() == 0:
            window_only = 0
            gtk.main()


class Clustmap(Hmap):
    """
    
    Clustmap
        
        Adds dendogram and labels to consensus matrix Hmap

        Usage:

            m = Clustmap(clust_data, labels = None, bsize = 10, tree_space = 200)

            m.save('filename')
            m.show()

            clust_data      - cluster.ConsensusCluster object which has a consensus_cluster attribute and optionally sample_id
                              labels for each sample in clust_data.datapoints
            labels          - Optional list of (sample label, cluster id) pairs.  It should be the same length as clust_data.datapoints and it should
                              be in the same order as clust_data.consensus_matrix.  This is important!  ConsensusCluster reorders
                              its matrix but ignores datapoints, because that is input from elsewhere.  
            bsize           - Size of heatmap datapoints.  Don't make this less than 10 unless you have a font which is smaller than 10pt.
            tree_space      - 200 is a good default.  If there is no tree (no reordering, or reordered by PAM method, for example),
                              this is ignored.

    """

    def __init__(self, clust_data, labels = None, bsize = 10, tree_space = 200, fontfile = '/usr/share/fonts/TTF/arial.ttf'):
        
        self.space = tree_space
        
        colours = ['red', 'blue', 'green', 'cyan', 'magenta', 'brown', 'orange']
        self.colour_map = self._init_colours(colours, [ x.cluster_id for x in clust_data.datapoints ])
        
        if labels is None:
            # (sample_id, cluster_id) pairs
            labels = [ (clust_data.datapoints[x].sample_id, clust_data.datapoints[x].cluster_id) for x in clust_data.reorder_indices ]

        # Try to load fontfile, and if that fails, load included backup
        # If THAT fails, load PIL builtin
        try:
            from PIL import ImageFont
            try:
                self.font = ImageFont.truetype(fontfile, 10)
            except:
                self.font = ImageFont.load('helvR08.pil') #Copyright (c) 1987 Adobe Systems, Inc., Portions Copyright 1988 Digital Equipment Corp.
        
        except IOError:
            self.font = None

        if len(clust_data.consensus_matrix) != len(labels):
            raise ValueError, "Number of columns and column label arrays have different lengths!"

        Hmap.__init__(self, clust_data.consensus_matrix, bsize = bsize) #Creates image in self.im if PIL is imported

        if self.im is not None:

            from PIL import Image, ImageDraw

            old_draw = ImageDraw.Draw(self.im)

            self.max_textsize = 0
            for label, clustid in labels:
                self.max_textsize = max(self.max_textsize, old_draw.textsize(label, font=self.font)[0])
            
            del old_draw #Keep GC from keeping the old image around

            if clust_data.tree is None:
                self.space = self.max_textsize + 5

            #Prepare
            newsize = (self.im.size[1] + self.space, self.im.size[0])  #To hold our rotated copy and some text
            im = Image.new('RGBA', newsize, 'white')
            
            #Trick to make vertical text when we're done, and add tree space
            im.paste(self.im.rotate(-90), (0, 0, self.im.size[1], self.im.size[0]))

            self.im = im
            self.draw = ImageDraw.Draw(self.im)
            
            #Actual work
            self._add_cluster_labels(labels)

            if clust_data.tree is not None:
                self._draw_dendogram(clust_data.tree)

            #Finish
            self.im = self.im.rotate(90)

    def _init_colours(self, colours, id_lst):
        """Assigns colours to a list of cluster_ids"""

        next_colour = cycle(colours)
        
        cluster_ids = dict.fromkeys(id_lst).keys()    

        col_map = dict(zip(cluster_ids, next_colour))
        col_map[None] = 'black'

        return col_map
        
    def _add_cluster_labels(self, labels):
        """Adds labels to the cluster image"""

        x_ptr = self.im.size[0] - self.space + 4
        y_ptr = 0

        for i in xrange(len(labels)):
            label, clustid = labels[i]
            self.draw.text((x_ptr, y_ptr + i*self.bsize), label, fill=self.colour_map[clustid], font=self.font)

    def _draw_dendogram(self, tree):
        """Draws the tree.  It's a bit terrifying, I'll admit."""

        tree_seq = tree.sequence

        x_ptr = self.im.size[0] - self.space + 6 + self.max_textsize
        y_ptr = self.bsize / 2

        x_step = (self.im.size[0] - x_ptr) / tree.depth

        x_loc = lambda scale: x_ptr + scale * x_step
        y_loc = lambda i: y_ptr + tree_seq.index(i) * self.bsize

        stack = [tree]
        active_level = [tree.depth]
        seen = []
        midpoints = []

        def draw_node(x_min_left, x_min_right, x_max, y_min, y_max, depth, colour):

            self.draw.line((x_min_left, y_min, x_max, y_min, x_max, y_min, x_max, y_max, x_min_right, y_max, x_max, y_max), fill=colour)
            midpoints.append(((abs(y_min - y_max) / 2 + min(y_min, y_max)), depth))

        def draw_bridge(endpoint):

            y_max = midpoints.pop()
            last_level = seen.pop()

            colour = self.colour_map[last_level.cluster_id]
            draw_node(x_loc(endpoint[1]), x_loc(y_max[1]), x_loc(last_level.depth), endpoint[0], y_max[0], last_level.depth, colour)

        while len(stack):
            next = stack.pop()
            level = active_level.pop()

            if seen and level > seen[-1].depth:
                stack.append(next)
                active_level.append(level)
                
                draw_bridge(midpoints.pop())
                continue

            if next.value is not None:
                draw_bridge([y_loc(next.value), 0])
            
            elif next.right.value is not None and next.left.value is not None:
                colour = self.colour_map[next.cluster_id]
                draw_node(x_ptr, x_ptr, x_loc(next.depth), y_loc(next.left.value), y_loc(next.right.value), next.depth, colour)

            else:
                seen.append(next)

                if next.right.value is None:
                    stack.extend([next.left, next.right])
                elif next.left.value is None:
                    stack.extend([next.right, next.left])

                active_level.extend([next.depth, next.depth])

        while len(seen):
            draw_bridge(midpoints.pop())


class Plot(object):

    """

    Plot

        Small front-end class to matplotlib's extensive plotting functions.

        Plot will take a list of 2 by X numpy matrices and a list of labels (optional) and draw them to the same grid
        with different (rotating) colours.  The labels, if set, will become a legend identifying these sets.

        Usage:

        Plot(plots, fig_label = '', legend = None)

            plots       - a list of numpy matrices to plot, each with 2 rows.  Row 0 is x-coord, row 1 is y-coord.

            fig_label   - The saved filename will become fig_label.png.  If fig_label is missing, it will be called "figure.png"

            legend      - a list of labels, same length as plots

        When Plot is initialised, the plot will be drawn and an image will be saved.

    """

    def __init__(self, plots, fig_label = None, legend = None):

        self.colours = cycle(['bo', 'go', 'ro', 'co', 'mo', 'ko', 'yo', 'bs', 'gs', 'rs', 'cs', 'ms', 'ks', 'ys', 'b^', 'g^', 'r^', 'c^', 'm^', 'k^', 'y^'])

        self.legend = legend

        if self.legend is not None:
            if len(plots) != len(legend):
                raise ValueError, "Plots and legend have different lengths!"

        self._write_fig(plots, fig_label)

    def _write_fig(self, plots, fig_label):
        
        from matplotlib.figure import Figure
        from matplotlib.backends.backend_ps import FigureCanvasBase as FigureCanvas

        fig = Figure()
        ax = fig.add_subplot(111)

        if self.legend is not None:
            plotlist = numpy.argsort(self.legend)
        else:
            plotlist = xrange(len(plots))

        for i in plotlist:

            if plots[i].shape[0] != 2:
                raise ValueError, "Attempting to plot matrix with row count other than 2"
            
            if self.legend is not None:
                ax.plot(plots[i][0], plots[i][1], self.colours.next(), label=unicode(self.legend[i], 'utf-8'))
            else:
                ax.plot(plots[i][0], plots[i][1], self.colours.next())

        if self.legend is not None:
            handles, labels = ax.get_legend_handles_labels()
            lgd = ax.legend(handles, labels, loc='upper center', bbox_to_anchor=(0.5,-0.1))

            canvas = FigureCanvas(fig)
            canvas.figure.savefig(new_filename(fig_label, 'figure', '.png'), bbox_extra_artists=(lgd,), bbox_inches='tight')

        else:
            canvas = FigureCanvas(fig)
            canvas.figure.savefig(new_filename(fig_label, 'figure', '.png'))


if __name__ == '__main__':

    matrix = []

    for i in range(32):
        matrix.append([])
        for j in range(32):
            if i & 1:
                matrix[i].append(float(-1 * i * j * 0.5))
            else:
                matrix[i].append(float(i * j * 0.5))

    Hmap(matrix, bsize = 15).show()
